import time
import av
from tools.log import logger

class AVStreamDecoder:
    def __init__(self, conf, retry_timeout=2, retry_max=3, on_frame=None, reconnect_interval=1):
        """
        conf: {"url": "rtsp://..."}
        retry_timeout: socket 超时秒数
        retry_max: 失败重试次数
        on_frame: 每解码一帧的回调 function(frame)
        reconnect_interval: 重连等待秒数
        """
        self.url = conf["url"]
        self.retry_timeout = retry_timeout
        self.retry_max = retry_max
        self.on_frame = on_frame
        self.reconnect_interval = reconnect_interval
        self.video_stream = None
        self.container =None
        self._stop_flag = False

    def run(self):
        """启动循环拉流"""
        while not self._stop_flag:
            if not self._connect():
                logger.warning(f"Reconnect in {self.reconnect_interval}s")
                time.sleep(self.reconnect_interval)
                continue

            try:
                for frame in self.container.decode(self.video_stream):
                    if self._stop_flag:
                        break
                    rgb = frame.to_rgb().to_ndarray()
                    if self.on_frame:
                        self.on_frame(rgb)
            except Exception as e:
                logger.error(f"Unexpected error during decode: {e}")
            finally:
                self._release()
                if not self._stop_flag:
                    logger.info(f"Will attempt reconnect in {self.reconnect_interval}s")
                    time.sleep(self.reconnect_interval)

    def stop(self):
        """停止拉流"""
        self._stop_flag = True
        self._release()

    def _connect(self):
        """建立 RTSP 连接，带重试"""
        for attempt in range(self.retry_max):
            try:
                container = av.open(
                    self.url,
                    timeout=self.retry_timeout,
                    options={
                        "rtsp_transport": "tcp",
                        # "fflags": "nobuffer",
                        "flags": "low_delay",
                        "analyzeduration": "0",
                        "probesize": "32",
                        "stimeout": str(5*1_000_000),  # 5秒超时
                        "framedrop": "1",
                        "max_delay": "2000000",
                    }
                )
                self.container=container
                # 获取视频流
                self.video_stream = container.streams.video[0]
                # # 强制使用 Rockchip MPP 硬件解码
                # self.video_stream.codec_context.options = {"hwaccel": "rkmpp"}  # PyAV/FFmpeg 有效
                logger.info(f"Connected to {self.url}")
                return True
            except Exception as e:
                logger.error(f"Failed to open stream (attempt {attempt + 1}/{self.retry_max}): {e}")
                self._release()
                time.sleep(self.reconnect_interval)
        return False

    def _release(self):
        """释放 container"""
        if self.container:
            try:
                self.container.close()
            except Exception as e:
                logger.error(f"Error closing container: {e}")
            finally:
                self.container = None
                logger.info(f"Disconnected from {self.url}")
